{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Content under Creative Commons Attribution license CC-BY 4.0, code under BSD 3-Clause License © 2018  by D. Koehn, notebook style sheet by L.A. Barba, N.C. Clementi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<link href=\"https://fonts.googleapis.com/css?family=Merriweather:300,300i,400,400i,700,700i,900,900i\" rel='stylesheet' >\n",
       "<link href=\"https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,300i,400,400i,700,700i\" rel='stylesheet' >\n",
       "<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' >\n",
       "<style>\n",
       "\n",
       "@font-face {\n",
       "    font-family: \"Computer Modern\";\n",
       "    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');\n",
       "}\n",
       "\n",
       "\n",
       "#notebook_panel { /* main background */\n",
       "    background: rgb(245,245,245);\n",
       "}\n",
       "\n",
       "div.cell { /* set cell width */\n",
       "    width: 800px;\n",
       "}\n",
       "\n",
       "div #notebook { /* centre the content */\n",
       "    background: #fff; /* white background for content */\n",
       "    width: 1000px;\n",
       "    margin: auto;\n",
       "    padding-left: 0em;\n",
       "}\n",
       "\n",
       "#notebook li { /* More space between bullet points */\n",
       "margin-top:0.5em;\n",
       "}\n",
       "\n",
       "/* draw border around running cells */\n",
       "div.cell.border-box-sizing.code_cell.running { \n",
       "    border: 1px solid #111;\n",
       "}\n",
       "\n",
       "/* Put a solid color box around each cell and its output, visually linking them*/\n",
       "div.cell.code_cell {\n",
       "    background-color: rgb(256,256,256); \n",
       "    border-radius: 0px; \n",
       "    padding: 0.5em;\n",
       "    margin-left:1em;\n",
       "    margin-top: 1em;\n",
       "}\n",
       "\n",
       "\n",
       "div.text_cell_render{\n",
       "    font-family: 'Source Sans Pro', sans-serif;\n",
       "    line-height: 140%;\n",
       "    font-size: 110%;\n",
       "    width:680px;\n",
       "    margin-left:auto;\n",
       "    margin-right:auto;\n",
       "}\n",
       "\n",
       "/* Formatting for header cells */\n",
       ".text_cell_render h1 {\n",
       "    font-family: 'Merriweather', serif;\n",
       "    font-style:regular;\n",
       "    font-weight: bold;    \n",
       "    font-size: 250%;\n",
       "    line-height: 100%;\n",
       "    color: #004065;\n",
       "    margin-bottom: 1em;\n",
       "    margin-top: 0.5em;\n",
       "    display: block;\n",
       "}\t\n",
       ".text_cell_render h2 {\n",
       "    font-family: 'Merriweather', serif;\n",
       "    font-weight: bold; \n",
       "    font-size: 180%;\n",
       "    line-height: 100%;\n",
       "    color: #0096d6;\n",
       "    margin-bottom: 0.5em;\n",
       "    margin-top: 0.5em;\n",
       "    display: block;\n",
       "}\t\n",
       "\n",
       ".text_cell_render h3 {\n",
       "    font-family: 'Merriweather', serif;\n",
       "\tfont-size: 150%;\n",
       "    margin-top:12px;\n",
       "    margin-bottom: 3px;\n",
       "    font-style: regular;\n",
       "    color: #008367;\n",
       "}\n",
       "\n",
       ".text_cell_render h4 {    /*Use this for captions*/\n",
       "    font-family: 'Merriweather', serif;\n",
       "    font-weight: 300; \n",
       "    font-size: 100%;\n",
       "    line-height: 120%;\n",
       "    text-align: left;\n",
       "    width:500px;\n",
       "    margin-top: 1em;\n",
       "    margin-bottom: 2em;\n",
       "    margin-left: 80pt;\n",
       "    font-style: regular;\n",
       "}\n",
       "\n",
       ".text_cell_render h5 {  /*Use this for small titles*/\n",
       "    font-family: 'Source Sans Pro', sans-serif;\n",
       "    font-weight: regular;\n",
       "    font-size: 130%;\n",
       "    color: #e31937;\n",
       "    font-style: italic;\n",
       "    margin-bottom: .5em;\n",
       "    margin-top: 1em;\n",
       "    display: block;\n",
       "}\n",
       "\n",
       ".text_cell_render h6 { /*use this for copyright note*/\n",
       "    font-family: 'Source Code Pro', sans-serif;\n",
       "    font-weight: 300;\n",
       "    font-size: 9pt;\n",
       "    line-height: 100%;\n",
       "    color: grey;\n",
       "    margin-bottom: 1px;\n",
       "    margin-top: 1px;\n",
       "}\n",
       "\n",
       "    .CodeMirror{\n",
       "            font-family: \"Source Code Pro\";\n",
       "\t\t\tfont-size: 90%;\n",
       "    }\n",
       "/*    .prompt{\n",
       "        display: None;\n",
       "    }*/\n",
       "\t\n",
       "    \n",
       "    .warning{\n",
       "        color: rgb( 240, 20, 20 )\n",
       "        }  \n",
       "</style>\n",
       "<script>\n",
       "    MathJax.Hub.Config({\n",
       "                        TeX: {\n",
       "                           extensions: [\"AMSmath.js\"], \n",
       "                           equationNumbers: { autoNumber: \"AMS\", useLabelIds: true}\n",
       "                           },\n",
       "                tex2jax: {\n",
       "                    inlineMath: [ ['$','$'], [\"\\\\(\",\"\\\\)\"] ],\n",
       "                    displayMath: [ ['$$','$$'], [\"\\\\[\",\"\\\\]\"] ]\n",
       "                },\n",
       "                displayAlign: 'center', // Change this to 'center' to center equations.\n",
       "                \"HTML-CSS\": {\n",
       "                    styles: {'.MathJax_Display': {\"margin\": 4}}\n",
       "                }\n",
       "        });\n",
       "</script>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Execute this cell to load the notebook's style sheet, then ignore it\n",
    "from IPython.core.display import HTML\n",
    "css_file = '../../style/custom.css'\n",
    "HTML(open(css_file, \"r\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computation of Sensitivity Kernels by 2D acoustic FD modelling\n",
    "\n",
    "Beside the modelling of seismic surveys, our 2D acoustic FD code can be used as the core of a seismic full waveform inversion (FWI) approach. A very efficient implementation is possible in the frequency domain.\n",
    "\n",
    "The aim of acoustic frequency domain FWI is to minimize the data residuals $\\mathbf{\\delta \\tilde{P} = \\tilde{P}^{mod} - \\tilde{P}^{obs}}$ between the modelled frequency domain data $\\mathbf{\\tilde{P}^{mod}}$ and field data $\\mathbf{\\tilde{P}^{obs}}$ to deduce high resolution models of the P-wave velocity distribution in the subsurface. To solve this non-linear inversion problem, an objective function E, as a measure of the data misfit, has to be defined. The classical choice is the L2-norm of the data residuals\n",
    "\n",
    "\\begin{equation} \n",
    "E  = ||\\mathbf{\\delta \\tilde{P}}||_2 = \\frac{1}{2}\\mathbf{\\delta \\tilde{P}}^\\dagger \\mathbf{\\delta \\tilde{P}} = \\frac{1}{2} \\sum_{k=1}^{n_\\omega} \\sum_{i=1}^{ns} \\sum_{j=1}^{nr} \\delta \\tilde{P}^*(\\mathbf{x_s}_i, \\mathbf{x_r}_j, \\omega_k) \\delta \\tilde{P}(\\mathbf{x_s}_i, \\mathbf{x_r}_j, \\omega_k) \\notag\n",
    "\\end{equation}  \n",
    "\n",
    "where ns and nr are the number of shots and receivers, $n_\\omega$ the number of discrete frequencies, $\\dagger$ the complex transpose, $*$ the complex conjugate, $\\mathbf{x_s},\\; \\mathbf{x_r}$  the source and receiver positions, respectively. The objective function can be minimized by iteratively updating the P-wave velocity $\\mathbf{Vp}$ at iteration step n, starting with an initial background model $\\mathbf{Vp_0}$, along a search direction using the **Newton** method:\n",
    "\n",
    "\\begin{equation} \n",
    "\\mathbf{Vp}_{n+1} = \\mathbf{Vp}_{n} - \\mu_n \\mathbf{H}_n^{-1} \\left(\\mathbf{\\frac{\\partial E}{\\partial Vp}}\\right)_n, \\notag\n",
    "\\end{equation}  \n",
    "\n",
    "where $\\mu$ denotes the step length, $\\mathbf{\\frac{\\partial E}{\\partial Vp}}$ the gradient and ${\\mathbf H}$ the second derivative (Hessian) of the objective function with respect to $\\mathbf{Vp}$. The step length $\\mu_{n}$ can be estimated by an inexact parabolic line search.\n",
    "\n",
    "The gradient $\\mathbf{\\frac{\\partial E}{\\partial Vp}}$ can be calculated by\n",
    "\n",
    "\\begin{equation} \n",
    "\\mathbf{\\frac{\\partial E}{\\partial Vp}} = - \\mathbf{K}^\\dagger \\delta P, \\notag\n",
    "\\end{equation}  \n",
    "\n",
    "where $\\mathbf{K}$ denotes the **Sensitivity Kernel**:\n",
    "\n",
    "\\begin{equation} \n",
    "K(x,z,\\omega) = {\\cal{Re}}\\biggl\\{\\frac{2 \\omega^2}{Vp^3} \\frac{\\tilde{P}(x,z,\\omega,\\mathbf{x_{s}}) \\tilde{G}(x,z,\\omega,\\mathbf{x_{r}})}{max(\\tilde{P}(x,z,\\omega,\\mathbf{x_{s}}))}\\biggr\\}\n",
    "\\end{equation} \n",
    "\n",
    "with the monochromatic forward wavefield $\\tilde{P}(x,z,\\omega,\\mathbf{x_s})$ excitated at the source position $\\mathbf{x_s}$ and the Green's function $\\tilde{G}(x,z,\\omega,\\mathbf{x_r})$ excitated at the receiver postion $\\mathbf{x_r}$, $\\cal{Re}$ denotes the real part.\n",
    "\n",
    "## Computation of monochromatic frequency domain wavefields from time-domain wavefields by Discrete Fourier Transform (DFT)\n",
    "\n",
    "To compute the sensitivity kernel eq. (1), we first need to estimate monochromatic frequency domain wavefields from the time domain wavefields. This can be easily implemented in our 2D acoustic FD code by the **Discrete Fourier Transform (DFT)** within the time-loop of the FD code. We approximate the continous Fourier transform \n",
    "\n",
    "\\begin{equation}\n",
    "\\tilde{f}(\\omega) = \\frac{1}{2 \\pi} \\int_{-\\infty}^{\\infty} f(t) exp(-i \\omega t) dt \\notag\n",
    "\\end{equation}\n",
    "\n",
    "by\n",
    "\n",
    "\\begin{equation}\n",
    "\\tilde{f_i}(\\omega) \\approx \\frac{1}{2 \\pi} \\sum_{n=0}^{nt} f_n(t_n) \\biggl(cos(\\omega t_n)-i\\; sin(\\omega t_n)\\biggr) dt \\notag\n",
    "\\end{equation}\n",
    "\n",
    "with $nt$ the number of time steps in the FD code, $\\omega = 2 \\pi f$ the circular frequency based on the frequency $f$, $i^2 = -1$. The wavefield can be further decomposed into the real \n",
    "\n",
    "\\begin{equation}\n",
    "Re\\{\\tilde{f_i}(\\omega)\\} \\approx \\frac{1}{2 \\pi} \\sum_{n=0}^{nt} f_n(t_n) \\biggl(cos(2 \\pi f t_n)\\biggr) dt \\notag\n",
    "\\end{equation}\n",
    "\n",
    "and imaginary part\n",
    "\n",
    "\\begin{equation}\n",
    "Im\\{\\tilde{f_i}(\\omega)\\} \\approx -\\frac{1}{2 \\pi} \\sum_{n=0}^{nt} f_n(t_n) \\biggl(sin(2 \\pi f t_n)\\biggr) dt \\notag\n",
    "\\end{equation}\n",
    "\n",
    "The implementation into the FD code is quite straightforward. During the time-stepping we have to multiply the time-domain pressure wavefield by a geometrical factor and add the contributions. Let's implement it into our 2D acoustic FD code and try to compute the frequency domain wavefields for a homogeneous background model and a first arrival traveltime tomography result of the Marmousi-2 model ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# Import Libraries \n",
    "# ----------------\n",
    "import numpy as np\n",
    "from numba import jit\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from pylab import rcParams\n",
    "\n",
    "# Ignore Warning Messages\n",
    "# -----------------------\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, we start with the definition of the basic modelling parameters ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Definition of modelling parameters\n",
    "# ----------------------------------\n",
    "\n",
    "# Define model discretization\n",
    "nx = 500    # number of grid points in x-direction\n",
    "nz = 174    # number of grid points in z-direction\n",
    "dx = 20.0   # spatial grid point distance in x-direction (m)\n",
    "dz = dx     # spatial grid point distance in z-direction (m)\n",
    "\n",
    "# Define xmax, zmax\n",
    "xmax = nx * dx\n",
    "zmax = nz * dz\n",
    "\n",
    "# Define maximum recording time\n",
    "tmax =  6.0 # maximum wave propagation time (s)\n",
    "\n",
    "# Define source and receiver position\n",
    "xsrc =  2000.0 # x-source position (m)\n",
    "zsrc =  40.0   # z-source position (m)\n",
    "\n",
    "xrec =  8000.0 # x-receiver position (m)\n",
    "zrec =  40.0   # z-source position (m)\n",
    "\n",
    "f0   = 10 # dominant frequency of the source (Hz)\n",
    "print(\"f0 = \", f0, \" Hz\")\n",
    "t0   = 4.0/f0   # source time shift (s)\n",
    "\n",
    "isnap = 2  # snapshot interval (timesteps)\n",
    "\n",
    "# Calculate monochromatic wavefields for discrete frequency freq\n",
    "freq = 5.0 # discrete frequency (Hz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... define a JIT-ed function for the spatial FD approximation ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True) # use JIT for C-performance\n",
    "def update_d2px_d2pz_3pt(p, dx, dz, nx, nz, d2px, d2pz):\n",
    "    \n",
    "    for i in range(1, nx - 1):\n",
    "        for j in range(1, nz - 1):\n",
    "                \n",
    "            d2px[i,j] = (p[i + 1,j] - 2 * p[i,j] + p[i - 1,j]) / dx**2                \n",
    "            d2pz[i,j] = (p[i,j + 1] - 2 * p[i,j] + p[i,j - 1]) / dz**2\n",
    "        \n",
    "    return d2px, d2pz "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... initialize the absorbing boundary frame ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define simple absorbing boundary frame based on wavefield damping \n",
    "# according to Cerjan et al., 1985, Geophysics, 50, 705-708\n",
    "def absorb(nx,nz):\n",
    "\n",
    "    FW = 60     # thickness of absorbing frame (gridpoints)    \n",
    "    a = 0.0053\n",
    "    \n",
    "    coeff = np.zeros(FW)\n",
    "    \n",
    "    # define coefficients in absorbing frame\n",
    "    for i in range(FW):    \n",
    "        coeff[i] = np.exp(-(a**2 * (FW-i)**2))\n",
    "\n",
    "    # initialize array of absorbing coefficients\n",
    "    absorb_coeff = np.ones((nx,nz))\n",
    "\n",
    "    # compute coefficients for left grid boundaries (x-direction)\n",
    "    zb=0 \n",
    "    for i in range(FW):\n",
    "        ze = nz - i - 1\n",
    "        for j in range(zb,ze):\n",
    "            absorb_coeff[i,j] = coeff[i]\n",
    "\n",
    "    # compute coefficients for right grid boundaries (x-direction)        \n",
    "    zb=0\n",
    "    for i in range(FW):\n",
    "        ii = nx - i - 1\n",
    "        ze = nz - i - 1\n",
    "        for j in range(zb,ze):\n",
    "            absorb_coeff[ii,j] = coeff[i]\n",
    "\n",
    "    # compute coefficients for bottom grid boundaries (z-direction)        \n",
    "    xb=0 \n",
    "    for j in range(FW):\n",
    "        jj = nz - j - 1\n",
    "        xb = j\n",
    "        xe = nx - j\n",
    "        for i in range(xb,xe):\n",
    "            absorb_coeff[i,jj] = coeff[j]\n",
    "\n",
    "    return absorb_coeff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the 2D FD acoustic modelling code, we implement the DFT of the time-domain wavefields, by initializing the real and imaginary parts of the pressure wavefields, calculate the trigonometric factors for the DFT within the time-loop, apply the DFT to the pressure wavefield `p` for the discrete frequency `freq` and finally return the real and imaginary parts of the frequency domain wavefields ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FD_2D_acoustic code with JIT optimization\n",
    "# -----------------------------------------\n",
    "def FD_2D_acoustic_JIT(vp,dt,dx,dz,f0,xsrc,zsrc,op,freq):        \n",
    "    \n",
    "    # calculate number of time steps nt \n",
    "    # ---------------------------------\n",
    "    nt = (int)(tmax/dt)\n",
    "    \n",
    "    # locate source on Cartesian FD grid\n",
    "    # ----------------------------------\n",
    "    isrc = (int)(xsrc/dx)  # source location in grid in x-direction\n",
    "    jsrc = (int)(zsrc/dz)  # source location in grid in x-direction    \n",
    "    \n",
    "    # Source time function (Gaussian)\n",
    "    # -------------------------------\n",
    "    src  = np.zeros(nt + 1)\n",
    "    time = np.linspace(0 * dt, nt * dt, nt)\n",
    "\n",
    "    # 1st derivative of Gaussian\n",
    "    src  = -2. * (time - t0) * (f0 ** 2) * (np.exp(- (f0 ** 2) * (time - t0) ** 2)) \n",
    "    \n",
    "    # define clip value: 0.1 * absolute maximum value of source wavelet\n",
    "    clip = 0.1 * max([np.abs(src.min()), np.abs(src.max())]) / (dx*dz) * dt**2\n",
    "    \n",
    "    # Define absorbing boundary frame\n",
    "    # -------------------------------    \n",
    "    absorb_coeff = absorb(nx,nz)\n",
    "    \n",
    "    # Define squared vp-model\n",
    "    # -----------------------        \n",
    "    vp2 = vp**2\n",
    "    \n",
    "    # Initialize empty pressure arrays\n",
    "    # --------------------------------\n",
    "    p    = np.zeros((nx,nz)) # p at time n (now)\n",
    "    pold = np.zeros((nx,nz)) # p at time n-1 (past)\n",
    "    pnew = np.zeros((nx,nz)) # p at time n+1 (present)\n",
    "    d2px = np.zeros((nx,nz)) # 2nd spatial x-derivative of p\n",
    "    d2pz = np.zeros((nx,nz)) # 2nd spatial z-derivative of p \n",
    "    \n",
    "    # INITIALIZE ARRAYS FOR REAL AND IMAGINARY PARTS OF MONOCHROMATIC WAVEFIELDS HERE! \n",
    "    # --------------------------------------------------------------------------------\n",
    "      # real part of the monochromatic wavefield \n",
    "      # imaginary part of the monochromatic wavefield \n",
    "        \n",
    "    # Initalize animation of pressure wavefield \n",
    "    # -----------------------------------------    \n",
    "    fig = plt.figure(figsize=(7,3))  # define figure size\n",
    "    extent = [0.0,xmax,zmax,0.0]     # define model extension\n",
    "    \n",
    "    # Plot Vp-model\n",
    "    image = plt.imshow((vp.T)/1000, cmap=plt.cm.gray, interpolation='nearest', \n",
    "                        extent=extent)    \n",
    "    \n",
    "    # Plot pressure wavefield movie\n",
    "    image1 = plt.imshow(p.T, animated=True, cmap=\"RdBu\", alpha=.75, extent=extent, \n",
    "                          interpolation='nearest', vmin=-clip, vmax=clip)    \n",
    "    plt.title('Pressure wavefield')\n",
    "    plt.xlabel('x [m]')\n",
    "    plt.ylabel('z [m]')\n",
    "           \n",
    "    plt.ion()    \n",
    "    plt.show(block=False)\n",
    "    \n",
    "    # Calculate Partial Derivatives\n",
    "    # -----------------------------\n",
    "    for it in range(nt):\n",
    "    \n",
    "        # FD approximation of spatial derivative by 3 point operator\n",
    "        if(op==3):\n",
    "            d2px, d2pz = update_d2px_d2pz_3pt(p, dx, dz, nx, nz, d2px, d2pz)\n",
    "\n",
    "        # Time Extrapolation\n",
    "        # ------------------\n",
    "        pnew = 2 * p - pold + vp2 * dt**2 * (d2px + d2pz)\n",
    "\n",
    "        # Add Source Term at isrc\n",
    "        # -----------------------\n",
    "        # Absolute pressure w.r.t analytical solution\n",
    "        pnew[isrc,jsrc] = pnew[isrc,jsrc] + src[it] / (dx * dz) * dt ** 2\n",
    "        \n",
    "        # Apply absorbing boundary frame\n",
    "        # ------------------------------\n",
    "        p *= absorb_coeff\n",
    "        pnew *= absorb_coeff\n",
    "        \n",
    "        # Remap Time Levels\n",
    "        # -----------------\n",
    "        pold, p = p, pnew\n",
    "        \n",
    "        # Calculate frequency domain wavefield at discrete frequency freq by DFT\n",
    "        # ----------------------------------------------------------------------\n",
    "           # time\n",
    "           # real part\n",
    "           # imaginary part\n",
    "        \n",
    "        # Estimate real and imaginary part of pressur wavefield p by DFT\n",
    "        # --------------------------------------------------------------\n",
    "\n",
    "    \n",
    "        # display pressure snapshots \n",
    "        if (it % isnap) == 0:            \n",
    "            image1.set_data(p.T)\n",
    "            fig.canvas.draw()\n",
    "    \n",
    "    # Finalize computation of DFT\n",
    "    \n",
    "    # Return real and imaginary parts of the monochromatic wavefield"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modelling monochromatic frequency domain wavefields for a homogeneous acoustic medium\n",
    "\n",
    "Now, everything is assembled to compute frequency domain wavefields. We only have to define the discrete frequency `freq` for which the monochromatic wavefields should be calculated. Let's start with a homogeneous model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run 2D acoustic FD modelling with 3-point spatial operater\n",
    "# ----------------------------------------------------------\n",
    "%matplotlib notebook\n",
    "op = 3  # define spatial FD operator (3-point) \n",
    "\n",
    "# define homogeneous model with vp = 2500 m/s\n",
    "\n",
    "# Define time step\n",
    "dt =  dx / (np.sqrt(2) * np.max(vp_hom))# time step (s)\n",
    "\n",
    "p_hom_re, p_hom_im = FD_2D_acoustic_JIT(vp_hom,dt,dx,dz,f0,xsrc,zsrc,op,freq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The time domain wavefields seem to be correct. Let's take a look at the frequency domain wavefield:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "# Plot real and imaginary parts of monochromatic wavefields\n",
    "clip_seis = 5e-10\n",
    "extent_seis = [0.0,xmax/1000,zmax/1000,0.0]\n",
    "\n",
    "ax = plt.subplot(211)\n",
    "plt.imshow(p_hom_re.T, cmap=plt.cm.RdBu, aspect=1, vmin=-clip_seis, \n",
    "                   vmax=clip_seis, extent=extent_seis)\n",
    "\n",
    "plt.title('Real part of monochromatic wavefield')\n",
    "#plt.xlabel('x [km]')\n",
    "ax.set_xticks([]) \n",
    "plt.ylabel('z [km]')\n",
    "\n",
    "\n",
    "plt.subplot(212)\n",
    "plt.imshow(p_hom_im.T, cmap=plt.cm.RdBu, aspect=1, vmin=-clip_seis, \n",
    "                   vmax=clip_seis, extent=extent_seis)\n",
    "plt.title('Imaginary part of monochromatic wavefield')\n",
    "plt.xlabel('x [km]')\n",
    "plt.ylabel('z [km]')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modelling monochromatic frequency domain wavefields for the Marmousi-2 FATT model\n",
    "\n",
    "In the next step, we calculate monochromatic wavefields for the first arrival traveltime tomography (FATT) result of the Marmousi-2 model, which could be an initial model for a subsequent FWI. First, we import the FATT model to Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import FATT result for Marmousi-2 Vp model\n",
    "# ------------------------------------------\n",
    "\n",
    "# Define model filename\n",
    "name_vp = \"../marmousi-2/marmousi_II_fatt.vp\"\n",
    "\n",
    "# Open file and write binary data to vp\n",
    "f = open(name_vp)\n",
    "data_type = np.dtype ('float32').newbyteorder ('<')\n",
    "vp_fatt = np.fromfile (f, dtype=data_type)\n",
    "\n",
    "# Reshape (1 x nx*nz) vector to (nx x nz) matrix \n",
    "vp_fatt = vp_fatt.reshape(nx,nz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot Marmousi-2 vp-model\n",
    "# ------------------------\n",
    "%matplotlib notebook\n",
    "extent = [0, xmax/1000, zmax/1000, 0]\n",
    "fig = plt.figure(figsize=(7,3))  # define figure size\n",
    "image = plt.imshow((vp_fatt.T)/1000, cmap=plt.cm.viridis, interpolation='nearest', \n",
    "                   extent=extent)\n",
    "\n",
    "\n",
    "cbar = plt.colorbar(aspect=12, pad=0.02)\n",
    "cbar.set_label('Vp [km/s]', labelpad=10)\n",
    "plt.title('Marmousi-2 FATT model')\n",
    "plt.xlabel('x [km]')\n",
    "plt.ylabel('z [km]')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and run the time-domain modelling code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run 2D acoustic FD modelling with 3-point spatial operater\n",
    "# ----------------------------------------------------------\n",
    "%matplotlib notebook\n",
    "op = 3  # define spatial FD operator (3-point) \n",
    "\n",
    "# Define time step by CFL criterion\n",
    "dt =  dx / (np.sqrt(2) * np.max(vp_fatt))# time step (s)\n",
    "\n",
    "p_fatt_re, p_fatt_im = FD_2D_acoustic_JIT(vp_fatt,dt,dx,dz,f0,xsrc,zsrc,op,freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "# Plot real and imaginary parts of monochromatic wavefields\n",
    "clip_seis = 5e-9\n",
    "extent_seis = [0.0,xmax/1000,zmax/1000,0.0]\n",
    "\n",
    "ax = plt.subplot(211)\n",
    "plt.imshow(p_fatt_re.T, cmap=plt.cm.RdBu, aspect=1, vmin=-clip_seis, \n",
    "                   vmax=clip_seis, extent=extent_seis)\n",
    "\n",
    "plt.title('Real part of monochromatic wavefield')\n",
    "#plt.xlabel('x [km]')\n",
    "ax.set_xticks([]) \n",
    "plt.ylabel('z [km]')\n",
    "\n",
    "\n",
    "plt.subplot(212)\n",
    "plt.imshow(p_fatt_im.T, cmap=plt.cm.RdBu, aspect=1, vmin=-clip_seis, \n",
    "                   vmax=clip_seis, extent=extent_seis)\n",
    "plt.title('Imaginary part of monochromatic wavefield')\n",
    "plt.xlabel('x [km]')\n",
    "plt.ylabel('z [km]')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sensitivity Kernels\n",
    "\n",
    "The monochromatic frequency domain wavefields are the key to calculate sensitivty kernels, which will be computed in the next exercise ...\n",
    "\n",
    "##### Exercise\n",
    "\n",
    "Compute 5 Hz frequency domain sensitivity kernels \n",
    "\n",
    "\\begin{equation} \n",
    "K(x,z,\\omega) = {\\cal{Re}}\\biggl\\{\\frac{2 \\omega^2}{Vp^3} \\frac{\\tilde{P}(x,z,\\omega,\\mathbf{x_{s}}) \\tilde{G}(x,z,\\omega,\\mathbf{x_{r}})}{max(\\tilde{P}(x,z,\\omega,\\mathbf{x_{s}}))}\\biggr\\} \\notag\n",
    "\\end{equation}\n",
    "\n",
    "for the homogenous and FATT Marmousi-2 model. \n",
    "\n",
    "- The frequency domain forward wavefields $\\tilde{P}(x,z,\\omega,\\mathbf{x_{s}})$ for a source at $x_{s} = 2000.0\\; m$ and  $z_{s} = 40.0\\; m$ where already computated in the previous sections of the notebook (`p_hom_re`, `p_hom_im`,`p_fatt_re`, `p_fatt_im`). You only have to compute the receiver Green's functions $\\tilde{G}(x,z,\\omega,\\mathbf{x_{r}})$ by placing a source at the receiver position $x_{r} = 8000.0\\; m$ and  $z_{r} = 40.0\\; m$\n",
    "- Compute the sensitivity kernels $K(x,z,\\omega)$. Hint: In Python complex numbers are defined as `real_part + 1j*imag_part`. This can also be applied to `NumPy` arrays.\n",
    "- Plot, describe and interpret the sensitivity kernels for the homogeneous and Marmousi-2 FATT model. Where would you expect model updates in a subsequent FWI? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPUTE RECEIVER GREEN'S FUNCTION FOR HOMOGENEOUS MODEL HERE!\n",
    "# -------------------------------------------------------------\n",
    "%matplotlib notebook\n",
    "op = 3  # define spatial FD operator (3-point) \n",
    "\n",
    "# Define time step\n",
    "dt =  dx / (np.sqrt(2) * np.max(vp_hom))# time step (s)\n",
    "\n",
    "g_hom_re, g_hom_im = FD_2D_acoustic_JIT(vp_hom,dt,dx,dz,f0,xsrc,zsrc,op,freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "# COMPUTE SENSITVITY KERNEL FOR HOMOGENEOUS MODEL HERE!\n",
    "\n",
    "clip = 4e-18\n",
    "extent = [0.0,xmax/1000,zmax/1000,0.0] # define model extension\n",
    "\n",
    "fig = plt.figure(figsize=(7,3))  # define figure size\n",
    "\n",
    "# Plot Vp-model\n",
    "image = plt.imshow((vp_hom.T)/1000, cmap=plt.cm.gray, interpolation='nearest', \n",
    "                    extent=extent)    \n",
    "\n",
    "# Plot Sensitivity Kernel\n",
    "image1 = plt.imshow(K_hom.T, cmap=\"RdBu\", alpha=.75, extent=extent, \n",
    "                      interpolation='nearest', vmin=-clip, vmax=clip)    \n",
    "plt.title('Sensitivity Kernel (homogeneous model)')\n",
    "plt.xlabel('x [km]')\n",
    "plt.ylabel('z [km]')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPUTE RECEIVER GREEN'S FUNCTION FOR MARMOUSI-2 FATT MODEL HERE!\n",
    "# -----------------------------------------------------------------\n",
    "%matplotlib notebook\n",
    "op = 3  # define spatial FD operator (3-point) \n",
    "\n",
    "# Define time step\n",
    "dt =  dx / (np.sqrt(2) * np.max(vp_fatt))# time step (s)\n",
    "\n",
    "g_fatt_re, g_fatt_im = FD_2D_acoustic_JIT(vp_fatt,dt,dx,dz,f0,xsrc,zsrc,op,freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "# COMPUTE SENSITVITY KERNEL FOR Marmousi-2 FATT MODEL HERE!\n",
    "\n",
    "clip = 8e-17\n",
    "extent = [0.0,xmax/1000,zmax/1000,0.0] # define model extension\n",
    "\n",
    "fig = plt.figure(figsize=(7,3))  # define figure size\n",
    "\n",
    "# Plot Vp-model\n",
    "image = plt.imshow((vp_fatt.T)/1000, cmap=plt.cm.gray, interpolation='nearest', \n",
    "                    extent=extent)    \n",
    "\n",
    "# Plot Sensitivity Kernel\n",
    "image1 = plt.imshow(K_hom.T, cmap=\"RdBu\", alpha=.5, extent=extent, \n",
    "                      interpolation='nearest', vmin=-clip, vmax=clip)    \n",
    "plt.title('Sensitivity Kernel (Marmousi-2 FATT model)')\n",
    "plt.xlabel('x [km]')\n",
    "plt.ylabel('z [km]')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What we learned:\n",
    "\n",
    "- How to compute monochromatic frequency domain wavefields from time-domain wavefields by discrete Fourier transform (DFT)\n",
    "- Computation of sensitivity kernels"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
